using DataFrames
using QuadGK
using Distributions
using JLD
using Optim
using ForwardDiff
using CSV
using LinearAlgebra
using SpecialFunctions
using Random
using DelimitedFiles
using SparseArrays

clearconsole()

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
include(readDir *"vech.jl");
include(readDir *"VAR_Procedures.jl");
include(readDir *"Loaddata.jl");
include(readDir *"VAR_MoreLags_Procedures.jl")


#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nDataSel  = "3"
nfVARSpec = "1"
nVARSpec  = "6"
nVARMCMCSpec= "1"
specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/fVARspec" * nfVARSpec * ".jl")
include(specDir * "/VARspec" * nVARSpec * ".jl")
include(specDir * "/VARMCMCspec" * nVARMCMCSpec * ".jl")


#-------------------------------------------------------------
# load aggregate data (Zt and Kexact)
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/Para" * nDataSel *"/"
agg_data = loadaggdata(1,Tend)
n_agg    = size(agg_data)[2]


#-------------------------------------------------------------
# Load coefficients from density estimation
#-------------------------------------------------------------
sNameLoadDir = "fVAR" * nfVARSpec
loaddir      = "$(pwd())/results/Para" *nDataSel * "/" * sNameLoadDir *"/";
sNameFile    = "K" * string(K) * "_fVAR" * nfVARSpec

PhatDensCoef_factor, MDD_GoF, VinvLam_all, ~ , ~  = loaddensdata(1,Tend,K,nfVARSpec)
n_cross = size(PhatDensCoef_factor)[2]

#-------------------------------------------------------------
# Adjust sample
#-------------------------------------------------------------

# Arrange data for VAR analysis
# We start by dropping the initial Tdrop-nlags observations
agg_data_adj = agg_data[Tdrop-nlags+1:end,:]
PhatDensCoef_factor_adj = PhatDensCoef_factor[Tdrop-nlags+1:end,:]
# We now have T*=(T-T0)+nlags observations and call the first observation t=1
# Observations t=1,...,nlags are used to initialize lags
# and observations t=nlags+1,...,T* are the estimation sample
# Note that by holding Tdrop fixed and changing nlags<=Tdrop we can change
# the nunber of lags holding the estimation sample fixed.
# Estimation sample has T*-nlags=T-T0 observations

# Define some constants
n       = n_agg+n_cross
p       = nlags

# Combine aggregate and cross sectional observations and create YY XX for VAR estimation
data_adj = [ agg_data_adj PhatDensCoef_factor_adj ]
YY, XX, PHIols, Sols, SIGMAols = data_to_YY(data_adj, nlags, nlags)
# FS: having nlags twice as input arguments is not necessary. We can simplify
# but would have to adjust other programs as well.

#-------------------------------------------------------------
# Generate prior
#-------------------------------------------------------------
igammadf = size(YY)[2]+sigdf
# Test
# prior    = prior_ACP_redu(YY, SIGMAols, n_agg, igammadf, nlags, lambda)
prior    = prior_ACP_stru(YY, SIGMAols, n_agg, igammadf, nlags, lambda)

#-------------------------------------------------------------
# Run Posterior Sampler
#-------------------------------------------------------------

println("")
println("Bayesian estimation of VAR: Gibbs Sampling... ")
println("")

store_alp, store_beta, store_Sig = sample_ThetaSig(YY,XX,PHIols,Sols,SIGMAols,nlags,prior,nsim)
store_Btilde, store_Sigpdraw     = getReducedForm(store_alp,store_beta,store_Sig)

# Initialize matrices for collecting draws from posterior
store_Bpdraw     = zeros(nsim,n*p,n)
store_Sigtrpdraw = zeros(nsim,n,n)

for jj = 1:nsim

    Sigpdraw_j = Symmetric(store_Sigpdraw[jj,:,:])
    L = cholesky(Sigpdraw_j).L
    store_Sigtrpdraw[jj,:,:] = L

    Btilde = reshape(store_Btilde[jj,:], n*p,n)
    global store_Bpdraw[jj,:,:] = Btilde

end

store_Bpdraw     = store_Bpdraw[nburn+1:nsim,:,:];
store_Sigpdraw   = store_Sigpdraw[nburn+1:nsim,:,:];
store_Sigtrpdraw = store_Sigtrpdraw[nburn+1:nsim,:,:];


#-------------------------------------------------------------
# Save draws
#-------------------------------------------------------------

sName    = "fVAR" * nfVARSpec * "_VAR" * nVARSpec*"_MCMC" * nVARMCMCSpec
savedir  = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";
try mkdir(savedir) catch; end
save(savedir * sName* "_PostDraws.jld", "Bpdraw", store_Bpdraw, "Sigpdraw", store_Sigpdraw, "Sigtrpdraw", store_Sigtrpdraw)

#-------------------------------------------------------------
# MISC ITEMS
#-------------------------------------------------------------

Bpmean       = dropdims(mean(store_Bpdraw,dims=1),dims=1)
Sigpmean     = dropdims(mean(store_Sigpdraw,dims=1),dims=1)
Sigtrpmean   = dropdims(mean(store_Sigtrpdraw,dims=1),dims=1)

save(savedir * sName* "_PostMeans.jld", "Bpmean", Bpmean, "Sigpmean", Sigpmean, "Sigtrpmean", Sigtrpmean)

# save OLS estiamtes and posterior means as CSV files
CSV.write(savedir * sName * "_PHIols.csv", DataFrame(PHIols,:auto))
CSV.write(savedir * sName * "_SIGMAols.csv", DataFrame(SIGMAols,:auto))
CSV.write(savedir * sName * "_Bpmean.csv", DataFrame(Bpmean,:auto))
CSV.write(savedir * sName * "_Sigpmean.csv", DataFrame(Sigpmean,:auto))
CSV.write(savedir * sName * "_Sigtrpmean.csv", DataFrame(Sigtrpmean,:auto))
